from fastapi import FastAPI, UploadFile, File, HTTPException
from contextlib import asynccontextmanager
from fastapi.responses import JSONResponse
from funasr import AutoModel
from funasr.utils.postprocess_utils import rich_transcription_postprocess
import tempfile
import os
import shutil
import uuid

from starlette.middleware.cors import CORSMiddleware
from starlette.websockets import WebSocketDisconnect, WebSocket
import requests
import random
from hashlib import md5
from fastapi import Body

'''
启动命令   
uvicorn main:app --reload
'''
# 全局模型实例
model = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    """服务启动时初始化模型"""
    global model
    model_dir = "model/iic/SenseVoiceSmall"
    model = AutoModel(
        model=model_dir,
        vad_model="model/iic/speech_fsmn_vad_zh-cn-16k-common-pytorch",
        vad_kwargs={"max_single_segment_time": 30000},
        device="cuda:1",
        disable_update=True
    )
    yield
    # 清理资源
    if model:
        del model


app = FastAPI(lifespan=lifespan)

# 允许跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.post("/recognize")
async def recognize_speech(
        file: UploadFile = File(...),
        language: str = "auto",
        use_itn: bool = True,
        batch_size_s: int = 60,
        merge_vad: bool = True,
        merge_length_s: int = 15
):
    """语音识别接口"""
    temp_dir = None
    temp_path = None

    try:
        # 创建临时目录
        temp_dir = tempfile.mkdtemp()
        temp_path = os.path.join(temp_dir, f"upload_{uuid.uuid4().hex}{os.path.splitext(file.filename)[1]}")

        # 保存上传文件
        with open(temp_path, "wb") as f:
            shutil.copyfileobj(file.file, f)

        # 模型推理
        res = model.generate(
            input=temp_path,
            cache={},
            language=language,
            use_itn=use_itn,
            batch_size_s=batch_size_s,
            merge_vad=merge_vad,
            merge_length_s=merge_length_s
        )

        # 后处理
        processed_text = rich_transcription_postprocess(res[0]["text"])
        return {"text": processed_text}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        # 确保路径变量已定义
        if temp_path and os.path.exists(temp_path):
            try:
                os.remove(temp_path)
            except Exception as e:
                print(f"文件删除失败: {str(e)}")

        if temp_dir and os.path.exists(temp_dir):
            try:
                shutil.rmtree(temp_dir)  # 使用rmtree替代rmdir以删除非空目录
            except Exception as e:
                print(f"目录删除失败: {str(e)}")


# WebSocket实时识别接口
# 后端修改：接收二进制数据流
@app.websocket("/ws/stream-recognize")
async def websocket_stream_recognize(websocket: WebSocket):
    await websocket.accept()
    session_id = str(uuid.uuid4())
    temp_dir = None

    try:
        while True:
            # 接收二进制数据
            data = await websocket.receive_bytes()

            # 创建临时文件
            temp_dir = tempfile.mkdtemp()
            temp_path = os.path.join(temp_dir, f"stream_{session_id}.pcm")

            # 直接写入原始PCM数据
            with open(temp_path, "wb") as f:
                f.write(data)

            # 指定音频格式参数
            res = model.generate(
                input=temp_path,
                stream=True,
                cache={},
                language="auto",
                audio_format="pcm",
                sample_rate=16000
            )

            # 返回部分结果
            partial_text = rich_transcription_postprocess(res[0]["text"])
            print(partial_text)
            await websocket.send_json({"text": partial_text})

    except WebSocketDisconnect:
        print("客户端断开连接")
    finally:
        if temp_dir and os.path.exists(temp_dir):
            shutil.rmtree(temp_dir, ignore_errors=True)



# 百度翻译API配置（用环境变量）
BAIDU_APPID = os.getenv('BaiDuAPPID')
BAIDU_APPKEY = os.getenv('BaiDuAPPKEY')

def make_md5(s, encoding='utf-8'):
    return md5(s.encode(encoding)).hexdigest()

@app.post("/translate")
async def translate(
    text: str = Body(..., embed=True),
    to_lang: str = Body("en", embed=True)  # 目标语言参数，默认英文
):
    """
    多语言自动识别翻译接口，支持多种目标语言
    """
    if not BAIDU_APPID or not BAIDU_APPKEY:
        raise HTTPException(status_code=500, detail="Baidu翻译API密钥未配置")

    # 支持的目标语言列表（百度API官方代码）
    supported_langs = {
        'zh': '中文',      # Chinese
        'en': '英文',      # English
        'jp': '日文',      # Japanese
        'kor': '韩文',     # Korean
        'fra': '法文',     # French
        'de': '德文',      # German
        'spa': '西班牙文', # Spanish
        'ru': '俄文',      # Russian
        'it': '意大利文',  # Italian
        'pt': '葡萄牙文',  # Portuguese
        'ar': '阿拉伯文',  # Arabic
        'th': '泰文',      # Thai
        'tr': '土耳其文',  # Turkish
        'vie': '越南文',   # Vietnamese
        'id': '印尼文',    # Indonesian
        'ms': '马来文',    # Malay
        'hi': '印地文',    # Hindi
    }
    if to_lang not in supported_langs:
        raise HTTPException(status_code=400, detail=f"不支持的目标语言: {to_lang}")

    # 百度API支持自动识别源语言
    from_lang = 'auto'

    endpoint = 'http://api.fanyi.baidu.com'
    path = '/api/trans/vip/translate'
    url = endpoint + path

    salt = random.randint(32768, 65536)
    sign = make_md5(BAIDU_APPID + text + str(salt) + BAIDU_APPKEY)

    payload = {
        'appid': BAIDU_APPID,
        'q': text,
        'from': from_lang,
        'to': to_lang,
        'salt': salt,
        'sign': sign
    }
    headers = {'Content-Type': 'application/x-www-form-urlencoded'}

    try:
        r = requests.post(url, params=payload, headers=headers, timeout=8)
        result = r.json()
        print("百度翻译API返回：", result)
        if 'trans_result' in result:
            # 百度返回的 detected_src_lang 字段为自动识别的源语言
            detected_from = result.get('from', from_lang)
            return {
                "translated": result['trans_result'][0]['dst'],
                "from_lang": detected_from,
                "to_lang": to_lang
            }
        else:
            raise HTTPException(status_code=500, detail=str(result))
    except Exception as e:
        import traceback
        print("翻译异常：", traceback.format_exc())
        raise HTTPException(status_code=500, detail=str(e))